import os
import sys
import requests
from tqdm import tqdm

"""
下载SAM2.1模型权重的Python脚本
"""

# SAM 2.1模型权重URL
sam2p1_urls = {
    "tiny": "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_tiny.pt",
    "small": "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_small.pt",
    "base_plus": "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_base_plus.pt",
    "large": "https://dl.fbaipublicfiles.com/segment_anything_2/092824/sam2.1_hiera_large.pt"
}

# 下载目录
download_dir = "../sam2/checkpoints"


def download_file(url, save_path):
    """
    下载文件并显示进度条
    """
    try:
        # 发送GET请求
        response = requests.get(url, stream=True)
        response.raise_for_status()  # 检查请求是否成功
        
        # 获取文件大小
        total_size = int(response.headers.get('content-length', 0))
        
        # 显示下载进度条
        with open(save_path, 'wb') as file, tqdm(
            desc=os.path.basename(save_path),
            total=total_size,
            unit='iB',
            unit_scale=True,
            unit_divisor=1024,
        ) as bar:
            for data in response.iter_content(chunk_size=8192):
                size = file.write(data)
                bar.update(size)
        
        print(f"成功下载: {save_path}")
        return True
        
    except Exception as e:
        print(f"下载失败: {e}")
        if os.path.exists(save_path):
            os.remove(save_path)  # 清理部分下载的文件
        return False


def main():
    """
    主函数
    """
    print("====================================")
    print("       SAM2.1模型权重下载工具        ")
    print("====================================")
    
    # 创建下载目录
    os.makedirs(download_dir, exist_ok=True)
    
    # 询问用户要下载哪个模型
    print("请选择要下载的模型 (输入编号):")
    print("1. sam2.1_hiera_tiny.pt (最小，适合资源有限的环境)")
    print("2. sam2.1_hiera_small.pt (较小)")
    print("3. sam2.1_hiera_base_plus.pt (中等)")
    print("4. sam2.1_hiera_large.pt (大型，推荐，性能最好)")
    print("5. 下载所有模型")
    
    choice = input("请输入选择 [1-5]: ")
    
    models_to_download = []
    
    if choice == "1":
        models_to_download = [("large", sam2p1_urls["large"])]  # 默认为large模型
    elif choice == "2":
        models_to_download = [("small", sam2p1_urls["small"])]
    elif choice == "3":
        models_to_download = [("base_plus", sam2p1_urls["base_plus"])]
    elif choice == "4":
        models_to_download = [("large", sam2p1_urls["large"])]
    elif choice == "5":
        models_to_download = list(sam2p1_urls.items())
    else:
        print("无效的选择，默认下载large模型")
        models_to_download = [("large", sam2p1_urls["large"])]
    
    # 开始下载
    for model_name, url in models_to_download:
        filename = f"sam2.1_hiera_{model_name if model_name != 'base_plus' else 'base_plus'}.pt"
        save_path = os.path.join(download_dir, filename)
        
        # 检查文件是否已存在
        if os.path.exists(save_path):
            overwrite = input(f"文件 {filename} 已存在，是否覆盖? (y/n): ")
            if overwrite.lower() != 'y':
                print(f"跳过下载 {filename}")
                continue
        
        print(f"开始下载 {filename} ({url})")
        if not download_file(url, save_path):
            print(f"下载 {filename} 失败，程序退出")
            sys.exit(1)
    
    print("\n====================================")
    print("       模型权重下载完成！           ")
    print("====================================")


if __name__ == "__main__":
    main()